#include "RUDPEx.h"
#ifndef min
#define min(a,b)            (((a) < (b)) ? (a) : (b))
#endif
#if _WIN32
#include <ws2tcpip.h>
#else
#include <arpa/inet.h>
#include <string.h>
#endif

RUDPEx::RUDPEx()
	:RUDP()
{

}

RUDPEx::~RUDPEx()
{
	Terminate();
}

bool RUDPEx::Initialize(const char* ip, uint16_t port)
{
	if (segment > BUFF_SIZE - PACKAGE_HEAD_SIZE)
	{
		ErrorLog("segment length (%d) > MAX_BUFF_SIZE %d", segment, BUFF_SIZE - PACKAGE_HEAD_SIZE);
		throw;
	}
	pkgID = RCTP::RandomSingle() * 10000;
	return RUDP::Initialize(ip, port);
}

void RUDPEx::Terminate()
{
	RUDP::Terminate();
	while (!packages.empty())
	{
		Package* pkg = packages.begin()->second;
		packages.erase(packages.begin()->first);
		delete pkg;
	}
}


void RUDPEx::Update(uint32_t timeout /* = 0 */)
{
	RUDP::Update(timeout);
	if (GetTimestamp() - lastCheckTimeoutTs >= 1000)
	{
		lastCheckTimeoutTs = GetTimestamp();
		for (std::map<uint64_t, Package*>::iterator iter = packages.begin(); iter != packages.end();)
		{
			Package* pkg = iter->second;
			if (GetTimestamp() - pkg->updatets > 10000)
			{
				std::map<uint64_t, Package*>::iterator re = iter;
				iter++;
				packages.erase(re);
				delete pkg;
				continue;
			}
			iter++;
		}
	}
}

void RUDPEx::SendTo(uint8_t* buffer, uint32_t len, sockaddr_in* remote, bool reliable)
{
	if (reliable && !isJoin)
	{
		uint64_t pid = AssemblyPackageID(pkgID++, RCTP::RandomSingle() * 10000, RCTP::GetTimestamp());
		int offset = 0;
		uint32_t msgLen = (uint32_t)len;
		uint16_t segID = 0;
		uint16_t segCount = msgLen / segment + (msgLen % segment == 0 ? 0 : 1);
		while (offset < msgLen)
		{
			uint32_t len = (uint32_t)min(msgLen - offset, segment);
			uint8_t* data = AssemblyPackageBuffer(pid, msgLen, segCount, segID, buffer + offset, len);
			RUDP::SendTo(data, len + PACKAGE_HEAD_SIZE, remote, true);
			delete[] data;
			offset += len;
			segID++;
		}
	}
	else
	{
		RUDP::SendTo(buffer, len, remote, false);
	}
}

void RUDPEx::SendTo(uint8_t* buffer, uint32_t len, const char* ip, const uint16_t port, bool reliable)
{
	RUDP::SendTo(buffer, len, ip, port, reliable);
}

void RUDPEx::OnRecvMsg(uint8_t* buffer, uint32_t len, bool reliable, UserData userdata)
{
	if (!reliable)
	{
		RUDP::OnRecvMsg(buffer, len, reliable, userdata);
		return;
	}
	if (!VerifyBuffer(buffer, len))
	{
		RUDP::OnRecvMsg(buffer, len, reliable, userdata);
		return;
	}
	uint64_t pid = *(uint64_t*)(buffer);
	uint32_t length = *(uint32_t*)(buffer + 8);
	uint16_t segCount = *(uint16_t*)(buffer + 12);
	uint16_t segID = *(uint16_t*)(buffer + 14);
	if (segCount == 1)
	{
		RUDP::OnRecvMsg(buffer + PACKAGE_HEAD_SIZE, len - PACKAGE_HEAD_SIZE, reliable, userdata);
	}
	else
	{
		auto result = packages.find(pid);
		if (result != packages.end())
		{
			Package* pkg = packages[pid];
			memcpy(pkg->data + segment * segID, buffer + PACKAGE_HEAD_SIZE, len - PACKAGE_HEAD_SIZE);
			pkg->crtLen += (uint32_t)(len - PACKAGE_HEAD_SIZE);
			pkg->updatets = GetTimestamp();
			if (pkg->crtLen >= pkg->length)
			{
				RUDP::OnRecvMsg(pkg->data, pkg->length, reliable, userdata);
				packages.erase(pid);
				delete pkg;
			}
		}
		else
		{
			Package* pkg = new Package();
			pkg->crtLen = (uint32_t)(len - PACKAGE_HEAD_SIZE);
			pkg->length = length;
			pkg->data = new uint8_t[length];
			pkg->updatets = GetTimestamp();
			memcpy(pkg->data + segment * segID, buffer + PACKAGE_HEAD_SIZE, len - PACKAGE_HEAD_SIZE);
			std::pair<std::map<uint64_t, Package*>::iterator, bool> insertPair;
			insertPair = packages.insert(std::map<uint64_t, Package*>::value_type(pid, pkg));
			if (!insertPair.second)
			{
				ErrorLog("Insert to packages failed");
				throw;
			}
		}
	}
}

uint64_t RUDPEx::AssemblyPackageID(uint16_t lid, uint16_t r, uint32_t ts)
{
	uint64_t value;
	uint8_t* data = (uint8_t*)&value;
	memcpy(data, &lid, 2);
	memcpy(data + 2, &r, 2);
	memcpy(data + 4, &ts, 4);
	return value;
}

uint8_t* RUDPEx::AssemblyPackageBuffer(uint64_t pid, uint32_t length, uint16_t segCount, uint16_t segID, uint8_t* data, uint32_t dataLen)
{
	uint8_t* buffer = new uint8_t[dataLen + PACKAGE_HEAD_SIZE];
	memcpy(buffer, &pid, 8);
	memcpy(buffer + 8, &length, 4);
	memcpy(buffer + 12, &segCount, 2);
	memcpy(buffer + 14, &segID, 2);
	uint16_t sum = 0;
	for (int i = 0; i < 16; i++)
		sum += buffer[i];
	memcpy(buffer + 16, &sum, 2);
	memcpy(buffer + 18, data, dataLen);
	return buffer;
}

bool RUDPEx::VerifyBuffer(uint8_t* buffer, uint32_t length)
{
	if (length < PACKAGE_HEAD_SIZE)
		return false;
	int sum = 0;
	for (int i = 0; i < 16; i++)
		sum += buffer[i];
	if (sum != *((uint16_t*)(buffer + 16)))
		return false;
	return true;
}

